import asyncio
import os
from typing import Any, Set, Callable, List
from loguru import logger
from ..config.live_config import LiveConfig

class TaskManager:
    def __init__(self) -> None:
        """初始化任务管理器"""
        self.task_urls: Set[str] = set()
        self.config = LiveConfig()
        self.running_tasks: Set[asyncio.Task] = set()  # 存储运行中的任务
        
    async def process_queue(self,
                          q: Any,
                          shared_list: list,
                          task_handler: Callable) -> bool:
        """处理队列中的任务"""
        try:
            # 检查可用槽位
            available_slots = self.config.MAX_TASK - len(self.task_urls)
            if available_slots <= 0:
                return False

            # 收集等待处理的任务
            batch = []
            while not q.empty() and len(batch) < min(available_slots, self.config.TASK_BATCH_SIZE):
                url = q.get()
                if url not in self.task_urls and url not in shared_list:
                    batch.append(url)

            if not batch:
                return True  # 没有更多任务需要处理

            # 处理当前批次
            
            # 并发创建所有任务
            tasks = []
            for url in batch:
                self.task_urls.add(url)
                shared_list.append(url)
                task = asyncio.create_task(task_handler(url))
                tasks.append(task)
                self.running_tasks.add(task)
                task.add_done_callback(lambda t: self.running_tasks.remove(t))
            
            # 等待所有任务至少启动
            await asyncio.sleep(0.1)
            
            # 输出当前状态
            # logger.info(f"当前任务数: {len(self.task_urls)}/{self.config.MAX_TASK}")
            return True

        except Exception as e:
            logger.error(f"处理队列任务出错: {str(e)}")
            return False
        
    def remove_task(self, url: str) -> None:
        """移除任务"""
        if url in self.task_urls:
            self.task_urls.remove(url)
            
    def has_capacity(self) -> bool:
        """检查是否还有任务容量"""
        return len(self.task_urls) < self.config.MAX_TASK